"""
策略执行引擎，负责策略的生命周期管理
"""
import asyncio
import logging
import pandas as pd
from typing import Dict, List, Type, Any, Optional
import importlib
import inspect
from datetime import datetime, timedelta

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, update, text

from models.strategy import Strategy as StrategyModel, Trade
from models.kline_5min import Kline5Min
from .base import StrategyBase

logger = logging.getLogger(__name__)

class StrategyEngine:
    """
    策略引擎，负责策略的注册、启动、停止和执行
    """
    
    def __init__(self, db_session: AsyncSession):
        """
        初始化策略引擎
        
        Args:
            db_session: 数据库会话对象
        """
        self.db_session = db_session
        self.registered_strategies: Dict[str, Type[StrategyBase]] = {}
        self.running_instances: Dict[str, StrategyBase] = {}
        self.tasks: Dict[str, asyncio.Task] = {}
    
    async def load_strategies_from_db(self):
        """从数据库加载策略配置"""
        query = select(StrategyModel).where(StrategyModel.is_active == True)
        result = await self.db_session.execute(query)
        strategies = result.scalars().all()
        
        for strategy_model in strategies:
            try:
                strategy_instance = await self._create_strategy_instance(
                    strategy_model.code,
                    strategy_model.name,
                    strategy_model.parameters
                )
                
                if strategy_instance:
                    instance_key = f"{strategy_model.code}_{strategy_model.parameters.get('symbol', 'default')}"
                    self.running_instances[instance_key] = strategy_instance
                    logger.info(f"策略 {strategy_model.name} 已从数据库加载")
            except Exception as e:
                logger.error(f"加载策略 {strategy_model.name} 失败: {str(e)}", exc_info=True)
    
    async def _create_strategy_instance(self, strategy_code: str, strategy_name: str, parameters: Dict[str, Any]) -> Optional[StrategyBase]:
        """创建策略实例"""
        if strategy_code not in self.registered_strategies:
            logger.warning(f"策略代码 {strategy_code} 未注册")
            return None
            
        try:
            strategy_class = self.registered_strategies[strategy_code]
            symbol = parameters.get("symbol", "BTC_USDT")  # 默认使用BTC_USDT
            return strategy_class(symbol, parameters)
        except Exception as e:
            logger.error(f"创建策略 {strategy_name} 实例失败: {str(e)}", exc_info=True)
            return None
    
    def register_strategy(self, strategy_class: Type[StrategyBase], code: str = None):
        """
        注册策略类
        
        Args:
            strategy_class: 策略类，必须继承自StrategyBase
            code: 策略代码，如果不提供则使用类名
        """
        if not issubclass(strategy_class, StrategyBase):
            raise TypeError("策略类必须继承自StrategyBase")
        
        strategy_code = code or strategy_class.__name__
        self.registered_strategies[strategy_code] = strategy_class
        logger.info(f"策略 {strategy_class.__name__} 已注册，代码: {strategy_code}")
    
    def register_strategies_from_module(self, module_name: str):
        """
        从模块中注册所有策略
        
        Args:
            module_name: 模块名，如 'services.strategies.implementations'
        """
        try:
            module = importlib.import_module(module_name)
            
            for name, obj in inspect.getmembers(module):
                if (inspect.isclass(obj) and 
                    issubclass(obj, StrategyBase) and 
                    obj != StrategyBase):
                    self.register_strategy(obj)
                    
            logger.info(f"从模块 {module_name} 注册了所有策略")
        except ImportError as e:
            logger.error(f"导入模块 {module_name} 失败: {str(e)}", exc_info=True)
    
    async def start_strategy(self, strategy_code: str, symbol: str, parameters: Dict[str, Any] = None) -> bool:
        """
        启动策略
        
        Args:
            strategy_code: 策略代码
            symbol: 交易对
            parameters: 策略参数
            
        Returns:
            bool: 是否成功启动
        """
        if strategy_code not in self.registered_strategies:
            logger.error(f"策略 {strategy_code} 未注册")
            return False
        
        instance_key = f"{strategy_code}_{symbol}"
        if instance_key in self.running_instances:
            logger.warning(f"策略 {strategy_code} 对于交易对 {symbol} 已经在运行")
            return False
        
        try:
            strategy_class = self.registered_strategies[strategy_code]
            strategy = strategy_class(symbol, parameters)
            self.running_instances[instance_key] = strategy
            
            # 启动策略执行任务
            task = asyncio.create_task(self._run_strategy_loop(strategy))
            self.tasks[instance_key] = task
            
            await strategy.on_start()
            logger.info(f"策略 {strategy_code} 对于交易对 {symbol} 已启动")
            return True
        except Exception as e:
            logger.error(f"启动策略 {strategy_code} 失败: {str(e)}", exc_info=True)
            return False
    
    async def stop_strategy(self, strategy_code: str, symbol: str) -> bool:
        """
        停止策略
        
        Args:
            strategy_code: 策略代码
            symbol: 交易对
            
        Returns:
            bool: 是否成功停止
        """
        instance_key = f"{strategy_code}_{symbol}"
        if instance_key not in self.running_instances:
            logger.warning(f"策略 {strategy_code} 对于交易对 {symbol} 未运行")
            return False
        
        try:
            strategy = self.running_instances[instance_key]
            await strategy.on_stop()
            
            # 取消任务
            if instance_key in self.tasks:
                task = self.tasks[instance_key]
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass
                del self.tasks[instance_key]
            
            del self.running_instances[instance_key]
            logger.info(f"策略 {strategy_code} 对于交易对 {symbol} 已停止")
            return True
        except Exception as e:
            logger.error(f"停止策略 {strategy_code} 失败: {str(e)}", exc_info=True)
            return False
    
    async def _run_strategy_loop(self, strategy: StrategyBase):
        """策略执行循环"""
        try:
            # 初始化策略状态
            await self._load_strategy_state(strategy)
            
            while True:
                try:
                    # 获取最新K线数据
                    klines = await self._fetch_recent_klines(strategy.symbol)
                    if klines.empty:
                        logger.warning(f"没有找到交易对 {strategy.symbol} 的K线数据")
                        await asyncio.sleep(60)  # 等待一分钟再试
                        continue
                    
                    # 准备数据
                    prepared_data = strategy.prepare_data(klines)
                    
                    # 注意：修改了执行顺序，先检查开仓信号，再检查平仓信号
                    
                    # 检查开仓信号
                    has_entry, entry_data = await strategy.check_entry_signal(prepared_data)
                    if has_entry:
                        logger.info(f"策略 {strategy.name} 对 {strategy.symbol} 产生开仓信号: {entry_data}")
                        await self._process_entry_signal(strategy, entry_data, prepared_data)
                    
                    # 检查当前持仓的平仓信号
                    open_trades = await self._get_open_trades(strategy)
                    for trade in open_trades:
                        entry_data = trade.signal_data
                        has_exit, exit_data = await strategy.check_exit_signal(prepared_data, entry_data)
                        
                        if has_exit:
                            logger.info(f"策略 {strategy.name} 对 {strategy.symbol} 产生平仓信号: {exit_data}")
                            await self._process_exit_signal(strategy, trade, exit_data, prepared_data)
                    
                    # 间隔检查
                    await asyncio.sleep(60)  # 每分钟检查一次
                    
                except asyncio.CancelledError:
                    raise
                except Exception as e:
                    logger.error(f"策略 {strategy.name} 执行出错: {str(e)}", exc_info=True)
                    await strategy.on_error(e)
                    await asyncio.sleep(300)  # 出错后休息5分钟再继续
                    
        except asyncio.CancelledError:
            logger.info(f"策略 {strategy.name} 任务被取消")
        except Exception as e:
            logger.error(f"策略 {strategy.name} 主循环异常: {str(e)}", exc_info=True)
            await strategy.on_error(e)
    
    async def _load_strategy_state(self, strategy: StrategyBase):
        """加载策略状态"""
        # 这里可以从数据库加载策略的当前状态
        pass
    
    async def _fetch_recent_klines(self, symbol: str) -> pd.DataFrame:
        """获取最近的K线数据"""
        # 获取过去100条5分钟K线
        query = select(Kline5Min).where(
            Kline5Min.symbol == symbol
        ).order_by(Kline5Min.datetime.desc()).limit(100)
        
        result = await self.db_session.execute(query)
        klines = result.scalars().all()
        
        if not klines:
            return pd.DataFrame()
        
        # 转换为DataFrame
        data = []
        for k in klines:
            data.append({
                'datetime': k.datetime,
                'datetime_utc': k.datetime_utc,
                'open': float(k.open),
                'high': float(k.high),
                'low': float(k.low),
                'close': float(k.close),
                'volume': float(k.volume),
                'amount': float(k.amount),
                'symbol': k.symbol
            })
            
        df = pd.DataFrame(data)
        # 按时间正序排列
        df = df.sort_values('datetime')
        return df
    
    async def _get_strategy_by_code(self, code: str) -> Optional[StrategyModel]:
        """通过代码获取策略模型"""
        query = select(StrategyModel).where(StrategyModel.code == code)
        result = await self.db_session.execute(query)
        return result.scalar_one_or_none()
    
    async def _get_open_trades(self, strategy: StrategyBase) -> List[Trade]:
        """获取当前开仓状态的交易"""
        strategy_model = await self._get_strategy_by_code(strategy.name)
        if not strategy_model:
            return []
            
        query = select(Trade).where(
            Trade.strategy_id == strategy_model.id,
            Trade.symbol == strategy.symbol,
            Trade.status == "open"
        )
        
        result = await self.db_session.execute(query)
        return result.scalars().all()
    
    async def _process_entry_signal(self, strategy: StrategyBase, entry_data: Dict[str, Any], market_data: pd.DataFrame):
        """处理开仓信号"""
        # 获取当前市场价格(假设用最后一条K线的收盘价)
        current_price = market_data['close'].iloc[-1]
        direction = entry_data.get('direction', 'long')  # 默认做多
        
        # 计算仓位大小
        position_size = strategy.calculate_position_size(current_price)
        
        # 获取策略数据库记录
        strategy_model = await self._get_strategy_by_code(strategy.name)
        if not strategy_model:
            logger.error(f"找不到策略 {strategy.name} 的数据库记录")
            return
        
        # 创建交易记录
        trade = Trade(
            strategy_id=strategy_model.id,
            symbol=strategy.symbol,
            direction=direction,
            entry_price=str(current_price),
            entry_time=datetime.utcnow(),
            volume=str(position_size),
            status="open",
            signal_data=entry_data
        )
        
        self.db_session.add(trade)
        await self.db_session.commit()
        logger.info(f"创建开仓交易记录: {trade}")
        
        # 这里可以添加实际执行交易的逻辑
    
    async def _process_exit_signal(self, strategy: StrategyBase, trade: Trade, exit_data: Dict[str, Any], market_data: pd.DataFrame):
        """处理平仓信号"""
        # 获取平仓价格，优先使用策略提供的价格，否则使用最后一条K线的收盘价
        exit_price = exit_data.get('price', market_data['close'].iloc[-1])
        
        # 计算盈亏
        entry_price = float(trade.entry_price)
        volume = float(trade.volume)
        
        profit = 0.0
        if trade.direction == "long":
            profit = (exit_price - entry_price) * volume
        else:  # short
            profit = (entry_price - exit_price) * volume
        
        # 更新交易记录
        trade.exit_price = str(exit_price)
        trade.exit_time = datetime.utcnow()
        trade.profit = str(profit)
        trade.status = "closed"
        
        await self.db_session.commit()
        logger.info(f"更新平仓交易记录: {trade}")
        
        # 这里可以添加实际执行交易的逻辑 